#pragma once
#include <iostream>
#include <string>
#include <direct.h> 
#include "dirent.h"
#include <vector>
#include <thread>
#include <chrono>
#include "file_utility.hpp"
#include "gpu_utility.hpp"
#include "logging.h"
#include "ennx_crypt.hpp"
#include <cuda_runtime_api.h>
#include "NvOnnxParser.h"

/**
 * @struct st_trans_model_input
 * @brief 模型转换所需的输入数据结构。
 */
struct st_trans_model_input
{
    std::string cfg_path;            // 配置文件路径
    std::vector<char> model_dat;     // 模型文件数据（已解密或原始数据）
    std::string gpu_model = "NVIDIA GeForce RTX 3060"; // GPU 型号
    std::string trans_path;          // 转换后的模型文件路径
};

// 用户定义的内存单位（如 1_GiB）
constexpr long long int operator"" _GiB(long long unsigned int val)
{
    return val * (1 << 30);
}
constexpr long long int operator"" _MiB(long long unsigned int val)
{
    return val * (1 << 20);
}
constexpr long long int operator"" _KiB(long long unsigned int val)
{
    return val * (1 << 10);
}

/**
 * @class ModelTransBase
 * @brief 模型转换基类，包含了模型的加载、解析、参数设置及引擎构建等功能。
 */
class ModelTransBase
{
public:
    // 构造函数
    ModelTransBase() {}

    // 析构函数，释放所有资源
    virtual ~ModelTransBase() { ResetAll(); }

    /**
     * @brief 转换模型，包括加载配置、解析 ONNX、设置参数和序列化引擎。
     *
     * @param input 模型转换所需的输入数据。
     * @return 转换成功返回 true，否则返回 false。
     */
    virtual bool TransModel(st_trans_model_input input)
    {
        m_input = input;

        // 依次调用各步骤，如果失败则返回 false
        if (!LoadCfg() ||
            !ParseOnnx() ||
            !SetParam() ||
            !BuildEngineSerialize())
        {
            return false;
        }

        ResetAll(); // 转换完成后重置所有资源
        return true;
    }

protected:
    // 加载配置的纯虚函数，需由子类实现
    virtual bool LoadCfg() = 0;

    /**
     * @brief 解析 ONNX 模型文件并构建 TensorRT 网络。
     *
     * @return 成功返回 true，否则返回 false。
     */
    virtual bool ParseOnnx()
    {
        ResetAll(); // 先重置资源

        // 创建 TensorRT builder
        builder = nvinfer1::createInferBuilder(sample::gLogger.getTRTLogger());
        assert(builder != nullptr);

        // 设置显式批处理
        const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
        network = builder->createNetworkV2(explicitBatch);
        config = builder->createBuilderConfig();

        // 使用 ONNX 解析器解析模型文件
        parser = nvonnxparser::createParser(*network, sample::gLogger.getTRTLogger());
        if (!parser->parse(m_input.model_dat.data(), m_input.model_dat.size()))
        {
            sample::gLogError << "Failure while parsing ONNX file" << std::endl;
            return false;
        }
        return true;
    }

    // 设置参数的纯虚函数，需由子类实现
    virtual bool SetParam() = 0;

    /**
     * @brief 构建 TensorRT 引擎并将其序列化。
     *
     * @return 成功返回 true，否则返回 false。
     */
    virtual bool BuildEngineSerialize()
    {
        std::cout << "start building engine" << std::endl;
        engine = builder->buildEngineWithConfig(*network, *config);
        std::cout << "build engine done" << std::endl;

        // 序列化引擎并保存数据
        nvinfer1::IHostMemory* data = engine->serialize();
        serialize_data.resize(data->size());
        memcpy(serialize_data.data(), data->data(), data->size());
        std::cout << "save engine file done" << std::endl;

        return true;
    }

    /**
     * @brief 重置所有资源，包括解析器、网络、构建器和引擎。
     */
    virtual void ResetAll()
    {
        // 销毁解析器
        if (parser)
            parser->destroy();
        parser = nullptr;

        // 销毁网络
        if (network)
            network->destroy();
        network = nullptr;

        // 销毁构建器
        if (builder)
            builder->destroy();
        builder = nullptr;

        // 销毁引擎
        if (engine)
            engine->destroy();
        engine = nullptr;
    }

protected:
    st_trans_model_input m_input; // 模型转换输入

    nvinfer1::ICudaEngine* engine = nullptr;      // TensorRT 引擎
    nvinfer1::IBuilder* builder = nullptr;        // TensorRT 构建器
    nvinfer1::INetworkDefinition* network = nullptr; // TensorRT 网络定义
    nvinfer1::IBuilderConfig* config = nullptr;   // TensorRT 配置
    nvonnxparser::IParser* parser = nullptr;      // ONNX 解析器

    std::vector<char> serialize_data;             // 序列化的引擎数据
};

/**
 * @brief 读取模型数据，根据扩展名进行解密或直接加载。
 *
 * @param model_path 模型文件路径。
 * @param v_dat 用于存储模型数据的向量。
 * @return 读取成功返回 true，否则返回 false。
 */
static bool read_model_dat(std::string model_path, std::vector<char>& v_dat)
{
    v_dat.clear();

    namespace fs = boost::filesystem;
    fs::path p(model_path.c_str());

    // 根据文件扩展名选择加载或解密
    if (p.extension().string() == ".onnx")
    {
        v_dat = get_file_data(model_path); // 直接读取 ONNX 文件数据
    }
    else if (p.extension().string() == ".ennx")
    {
        // 解密 ENNX 文件数据
        std::vector<char> v_e_model = get_file_data(model_path);
        std::string key = "irkdh_573%3?iq5h";
        std::string iv = "1234567890123456";
        v_dat = decryptEnnxFile(v_e_model, key, iv);
    }
    else
    {
        std::cout << "model_path extension no match : " << model_path.c_str() << std::endl;
        return false;
    }
    return !v_dat.empty();
}

/**
 * @brief 执行模型转换的主函数，遍历指定目录中的模型文件并转换为目标格式。
 *
 * @param model_trans 模型转换类实例。
 * @param mon_folder 监控的文件夹路径。
 * @param dest_folder 目标文件夹路径。
 * @param mark_info 标记信息。
 * @param ext 目标文件扩展名。
 */
static void models_trans(ModelTransBase& model_trans,
    std::string& mon_folder, std::string& dest_folder, std::string& mark_info, std::string ext)
{
    namespace fs = boost::filesystem;

    // 获取 GPU 型号
    std::string gpu_model = GetGPUModel();
    if (gpu_model == "")
    {
        std::cout << "Unknown Gpu Model! \n";
        return;
    }

    // 检查目标文件夹是否存在
    fs::path _dest_folder(dest_folder);
    if (!fs::exists(_dest_folder))
    {
        return;
    }

    // 获取监控文件夹中的所有配置文件
    std::vector<std::string> cfg_files = EnumFiles(mon_folder, std::string(".yaml"));
    for (int i = 0; i < cfg_files.size(); ++i)
    {
        st_trans_model_input input;
        fs::path _cfg_path(cfg_files[i]);
        input.cfg_path = _cfg_path.string();

        fs::path _model_path;

        // 查找与配置文件对应的模型文件
        if (fs::exists(_cfg_path.replace_extension(".ennx")))
        {
            _model_path = _cfg_path.replace_extension(".ennx");
        }
        else if (fs::exists(_cfg_path.replace_extension(".onnx")))
        {
            _model_path = _cfg_path.replace_extension(".onnx");
        }

        // 如果模型文件未找到，跳过此配置文件
        if (_model_path.empty())
        {
            std::cout << "no found ennx or onnx file !" << std::endl;
            continue;
        }

        // 设置目标路径
        fs::path dats_path(dest_folder + "\\" + _cfg_path.stem().string() + mark_info + ext);
        if (fs::exists(dats_path))
        {
            continue;
        }

        std::cout << "START===============================================\n";

        // 读取模型数据
        if (!read_model_dat(_model_path.string(), input.model_dat))
        {
            continue;
        }
        input.gpu_model = gpu_model;
        input.trans_path = dats_path.string() + "-tmp";

        // 转换模型
        if (model_trans.TransModel(input))
        {
            std::cout << "DONE================================================\n";
            int result = std::rename(input.trans_path.c_str(), dats_path.string().c_str());
            if (result == 0)
            {
                std::cout << "文件重命名成功\n";
            }
            else
            {
                std::cout << "文件重命名失败\n";
            }
        }
        else
        {
            std::cout << "FAIL!================================================\n";
        }
    }
}

/**
 * @brief 模型转换主程序入口，设置监控和目标文件夹路径，并调用模型转换函数。
 *
 * @param argc 参数数量。
 * @param argv 参数列表。
 * @param model_trans 模型转换类实例。
 * @param ext 目标文件扩展名（默认为 ".dats"）。
 */
static void model_trans_main(int argc, char* argv[], ModelTransBase& model_trans, std::string ext = ".dats")
{
    try
    {
        char cur_folder[260];
        _getcwd(cur_folder, sizeof(cur_folder)); // 获取当前工作目录
        std::string cwd = cur_folder;

        // 设置默认的监控和目标文件夹
        std::string mon_folder = cwd;
        std::string dest_folder = cwd;
        std::string mark_info = "";

        if (argc >= 2)
        {
            mon_folder = argv[1];
            dest_folder = mon_folder;
        }
        if (argc >= 3)
        {
            dest_folder = argv[2];
        }
        if (argc >= 4)
        {
            mark_info = argv[3];
        }

        // 执行模型转换
        models_trans(model_trans, mon_folder, dest_folder, mark_info, ext);
        printf("done!\n");
    }
    catch (...)
    {
        printf("exception!\n");
    }

    // 等待关闭
    for (int i = 0; i < 5; ++i)
    {
        std::this_thread::sleep_for(std::chrono::milliseconds(1000));
    }
}
